import pandas as pd
from datasets import Dataset
from setfit import (SetFitModel, Trainer, TrainingArguments,
                    get_templated_dataset)

labels = list(map(str.strip, open("./labels.txt").readlines()))

# A dummy dataset to fill with synthetic examples
dummy_dataset = Dataset.from_dict({})

train_dataset = get_templated_dataset(
    dummy_dataset,
    candidate_labels=labels,
    sample_size=2,
    template="反馈{}的问题。",
)

print("loading model")


# Initializing a new SetFit model
model = SetFitModel.from_pretrained("BAAI/bge-base-zh-v1.5", labels=labels)
print("load model done")

test_df = pd.read_excel("./data/test.xlsx")
test_df = test_df[["问题", "问题类别"]]
# rename column to "text" and "label"
test_df.columns = ["text", "label"]
test_df = test_df.dropna().applymap(str)
# filter label not in labels
test_df = test_df[test_df["label"].isin(labels)]
# transform label to index
test_df["label"] = test_df["label"].apply(lambda x: labels.index(x))
test_dataset = Dataset.from_dict(test_df.to_dict(orient="list"))

# Preparing the training arguments
args = TrainingArguments(
    batch_size=32,
    num_epochs=3,
)

# Preparing the trainer
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
)
trainer.train(
    eval_every=1,
    early_stopping_patience=3,
    early_stopping_metric="accuracy",
    load_best_model_at_end=True,
)

# Evaluating
metrics = trainer.evaluate(test_dataset)
print(metrics)
# => {'accuracy': 0.8511806699615596}

# Saving the trained model
model.save_pretrained("setfit-bge-base-v1.5-sst2-8-shot")
